// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/base/sorting_network.h>

#include <melon/benchmark.h>
#include <melon/random.h>
#include <melon/init/init.h>

#define POLLUX_BENCHMARK(_type, _name, ...) \
  [[maybe_unused]] _type _name(MELON_PP_STRINGIZE(_name), __VA_ARGS__)

namespace kumo::pollux {
namespace {

template <typename T>
class SortingNetworkBenchmark {
 public:
  SortingNetworkBenchmark(const char* name, int minLen, int maxLen) {
    int totalLen = 0;
    for (int i = 0; i < kNumSorts; ++i) {
      lengths_.push_back(melon::Random::rand32(minLen, maxLen));
      totalLen += lengths_.back();
    }
    data_.resize(totalLen);
    generateData();
    melon::addBenchmark(__FILE__, fmt::format("{}_std", name), [this] {
      return run([](auto* indices, int len, auto lt) {
        std::sort(indices, indices + len, lt);
      });
    });
    melon::addBenchmark(
        __FILE__, fmt::format("%{}_sorting_network", name), [this] {
          return run([](auto* indices, int len, auto lt) {
            sortingNetwork(indices, len, lt);
          });
        });
  }

 private:
  static constexpr int kNumSorts = 10'000;

  void generateData();

  template <typename Sort>
  unsigned run(Sort sort) const {
    std::vector<int32_t> indices;
    BENCHMARK_SUSPEND {
      for (int i = 0; i < kNumSorts; ++i) {
        indices.reserve(lengths_[i]);
      }
    }
    auto* buf = data_.data();
    for (int i = 0; i < kNumSorts; ++i) {
      indices.resize(lengths_[i]);
      std::iota(indices.begin(), indices.end(), 0);
      sort(indices.data(), lengths_[i], [&](auto i, auto j) {
        return buf[i] < buf[j];
      });
      buf += lengths_[i];
    }
    melon::doNotOptimizeAway(indices);
    return kNumSorts;
  }

  std::vector<int8_t> lengths_;
  std::vector<T> data_;
};

template <>
void SortingNetworkBenchmark<int32_t>::generateData() {
  for (int i = 0; i < data_.size(); ++i) {
    data_[i] = melon::Random::rand32();
  }
}

template <>
void SortingNetworkBenchmark<std::string>::generateData() {
  for (auto& s : data_) {
    s.resize(melon::Random::rand32(4, 32));
    for (int i = 0; i < s.size(); ++i) {
      s[i] = melon::Random::rand32(128);
    }
  }
}

struct ThreeWords {
  std::array<uint64_t, 3> value;

  bool operator<(const ThreeWords& other) const {
    return value < other.value;
  }
};

template <>
void SortingNetworkBenchmark<ThreeWords>::generateData() {
  for (auto& x : data_) {
    for (auto& y : x.value) {
      y = melon::Random::rand64();
    }
  }
}

} // namespace
} // namespace kumo::pollux

int main(int argc, char* argv[]) {
  using namespace kumo::pollux;
  melon::Init melonInit(&argc, &argv);
  POLLUX_BENCHMARK(SortingNetworkBenchmark<int32_t>, int32_2, 2, 4);
  POLLUX_BENCHMARK(SortingNetworkBenchmark<int32_t>, int32_4, 4, 8);
  POLLUX_BENCHMARK(SortingNetworkBenchmark<int32_t>, int32_8, 8, 16);
  POLLUX_BENCHMARK(SortingNetworkBenchmark<std::string>, string_2, 2, 4);
  POLLUX_BENCHMARK(SortingNetworkBenchmark<std::string>, string_4, 4, 8);
  POLLUX_BENCHMARK(SortingNetworkBenchmark<std::string>, string_8, 8, 16);
  POLLUX_BENCHMARK(SortingNetworkBenchmark<ThreeWords>, ThreeWords_2, 2, 4);
  POLLUX_BENCHMARK(SortingNetworkBenchmark<ThreeWords>, ThreeWords_4, 4, 8);
  POLLUX_BENCHMARK(SortingNetworkBenchmark<ThreeWords>, ThreeWords_8, 8, 16);
  melon::runBenchmarks();
  return 0;
}
